package com.sicau.chatmemory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.messages.Message;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

public class BaseOnFileChatMemoryRepository implements ChatMemoryRepository {
    private final String FILE_DIR;
    private static final Kryo kryo = new Kryo();

    static {
        kryo.setRegistrationRequired(false);
        kryo.setInstantiatorStrategy(new StdInstantiatorStrategy());
    }

    public BaseOnFileChatMemoryRepository(String FILE_DIR) {
        this.FILE_DIR = FILE_DIR;
        File file_dir = new File(FILE_DIR);
        if (!file_dir.exists()) {
            file_dir.mkdirs();
        }
    }

    @Override
    public List<String> findConversationIds() {
        return getConversationIds();
    }

    @Override
    public List<Message> findByConversationId(String conversationId) {
        List<Message> conversation = getConversation(conversationId);
        return conversation;
    }

    @Override
    public void saveAll(String conversationId, List<Message> messages) {
        List<String> conversationIds = getConversationIds();
        if (!conversationIds.contains(conversationId)) {
            conversationIds.add(conversationId);
            saveConversationIds(conversationIds);
        }
        saveConversation(conversationId, messages);
    }

    @Override
    public void deleteByConversationId(String conversationId) {
        List<String> conversationIds = getConversationIds();
        if (conversationIds.contains(conversationId)) {
            conversationIds.remove(conversationId);
            saveConversationIds(conversationIds);
        }
    }

    /**
     * 获取会话文件
     * @param conversationId
     * @return
     */
    private List<Message> getConversation(String conversationId) {
        File conversationFile = getConversationFile(conversationId);
        List<Message> messages = new ArrayList<>();
        if (conversationFile.exists()) {
            try(Input input = new Input(new FileInputStream(conversationFile));) {
                messages = kryo.readObject(input, ArrayList.class);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        return messages;
    }

    /**
     * 获取会话ids文件
     * @return
     */
    private List<String> getConversationIds() {
        File conversationFile = getConversationFile("ids");
        List<String> conversationIds = new ArrayList<>();
        if (conversationFile.exists()) {
            try(Input input = new Input(new FileInputStream(conversationFile));) {
                conversationIds = kryo.readObject(input, ArrayList.class);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        return conversationIds;
    }

    /**
     * 保存会话文件
     * @param conversationId
     * @param messages
     */
    private void saveConversation(String conversationId, List<Message> messages) {
        File conversationFile = getConversationFile(conversationId);
        try(Output output = new Output(new FileOutputStream(conversationFile))) {
            kryo.writeObject(output, messages);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 保存会话ids文件
     * @param conversationIds
     */
    private void saveConversationIds(List<String> conversationIds) {
        File conversationFile = getConversationFile("ids");
        try(Output output = new Output(new FileOutputStream(conversationFile))) {
            kryo.writeObject(output, conversationIds);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 获取会话文件夹
     * @return
     */
    private File getConversationFile(String conversationId) {
        return new File(FILE_DIR + "/conversationId_" + conversationId + ".kryo");
    }
}
